"""
model name : MY_YOLO
file       : loss.py
information:
    author : OuYang
    time   : 2025/1/22
"""

import torch
from torch import nn


class YOLOLoss(nn.Module):
    def __init__(
            self,
            coord=5,
            noobj=.5,
            s=7,
            b=2
    ):
        super(YOLOLoss, self).__init__()
        self.coord = coord
        self.noobj = noobj
        self.s = s
        self.b = b
        self.cell_size = 1.0 / self.s

    def forward(self, pred, target):
        batch_size = target.size(0)

        cls_loss = torch.tensor(0, dtype=torch.float32, device=target.device)
        conf_loss = torch.tensor(0, dtype=torch.float32, device=target.device)
        location_loss = torch.tensor(0, dtype=torch.float32, device=target.device)

        for batch_idx in range(batch_size):
            for i in range(self.s):
                for j in range(self.s):
                    # Grid[i][j] don't have object
                    if target[batch_idx][i][j][4] != 1:
                        conf_loss = conf_loss + self.noobj * torch.sum(
                            pred[batch_idx][i][j][[bi * 5 + 4 for bi in range(self.b)]] ** 2
                        )
                        continue

                    # Cls Loss
                    cls_loss = cls_loss + torch.sum(
                        (pred[batch_idx][i][j][self.b * 5:] - target[batch_idx][i][j][self.b * 5:]) ** 2
                    )

                    ious = []
                    max_iou = -1
                    max_bbox_idx = -1
                    for bbox_idx in range(self.b):
                        pred_x = pred[batch_idx][i][j][bbox_idx * 5]
                        pred_y = pred[batch_idx][i][j][bbox_idx * 5 + 1]
                        pred_w = pred[batch_idx][i][j][bbox_idx * 5 + 2]
                        pred_h = pred[batch_idx][i][j][bbox_idx * 5 + 3]
                        pred_box = (pred_x, pred_y, pred_w, pred_h)

                        target_x = target[batch_idx][i][j][bbox_idx * 5]
                        target_y = target[batch_idx][i][j][bbox_idx * 5 + 1]
                        target_w = target[batch_idx][i][j][bbox_idx * 5 + 2]
                        target_h = target[batch_idx][i][j][bbox_idx * 5 + 3]
                        target_box = (target_x, target_y, target_w, target_h)

                        iou = self._iou(pred_box, target_box, i, j)
                        ious.append(iou)

                        if iou > max_iou:
                            max_iou = iou
                            max_bbox_idx = bbox_idx

                    for bi, iou in enumerate(ious):
                        if bi == max_bbox_idx:
                            conf_loss = conf_loss + torch.sum((pred[batch_idx][i][j][max_bbox_idx * 5 + 4] - max_iou) ** 2)
                            location_loss = location_loss + self.coord * (
                                    torch.sum(
                                        (
                                                pred[batch_idx][i][j][bi * 5: bi * 5 + 2] -
                                                target[batch_idx][i][j][bi * 5:bi * 5 + 2]
                                        )
                                        ** 2
                                    )
                                    +
                                    torch.sum(
                                        (
                                                pred[batch_idx][i][j][bi * 5 + 2: bi * 5 + 4].sqrt() -
                                                target[batch_idx][i][j][bi * 5 + 2: bi * 5 + 4].sqrt()
                                        )
                                        ** 2
                                    )
                            )
                        else:
                            conf_loss = conf_loss + self.noobj * torch.sum(
                                pred[batch_idx][i][j][bi * 5 + 4] ** 2
                            )
        total_loss = cls_loss + conf_loss + location_loss
        return (total_loss / batch_size,
                cls_loss / batch_size,
                conf_loss / batch_size,
                location_loss / batch_size
                )

    def _iou(self, box1, box2, i, j):
        b1x, b1y, b1w, b1h = box1
        b2x, b2y, b2w, b2h = box2

        b1x1 = b1x * self.cell_size + i * self.cell_size - b1w / 2
        b1y1 = b1y * self.cell_size + j * self.cell_size - b1h / 2
        b1x2 = b1x * self.cell_size + i * self.cell_size + b1w / 2
        b1y2 = b1y * self.cell_size + j * self.cell_size + b1h / 2

        b2x1 = b2x * self.cell_size + i * self.cell_size - b2w / 2
        b2y1 = b2y * self.cell_size + j * self.cell_size - b2h / 2
        b2x2 = b2x * self.cell_size + i * self.cell_size + b2w / 2
        b2y2 = b2y * self.cell_size + j * self.cell_size + b2h / 2

        if b1x2 < b2x1 or b1x1 > b2x2 or b1y2 < b2y1 or b1y1 > b2y2:
            return 0
        inter_x1 = torch.max(b1x1, b2x1)
        inter_y1 = torch.max(b1y1, b2y1)
        inter_x2 = torch.min(b1x2, b2x2)
        inter_y2 = torch.min(b1y2, b2y2)

        b1_area = b1w * b1h
        b2_area = b2w * b2h
        inter_area = torch.abs(inter_x2 - inter_x1) * torch.abs(inter_y2 - inter_y1)

        union_area = b1_area + b2_area - inter_area
        iou = inter_area / torch.clamp(union_area, 1e-6)

        return iou if inter_area > 0 else 0

